import pandas as pd
import numpy as np
from collections import deque
from datetime import datetime


class THSNineTurn:
    """
    同花顺风格的九转策略（与主程实现等价）。
    """

    def __init__(self, compare_offset: int = 4, required_bars: int = 9,
                 max_window_minutes: int = 240, verbose: bool = True,
                 restrict_session: bool = True):
        self.compare_offset = int(compare_offset)
        self.required_bars = int(required_bars)
        self.max_window_minutes = int(max_window_minutes)
        self.verbose = bool(verbose)
        self.restrict_session = bool(restrict_session)

        # 状态存储
        self._bar_state = {}  # code -> {minute:int, open:float, high:float, low:float, close:float}
        self._closes = {}  # code -> deque of close prices
        self._counts = {}  # code -> {"up":int, "down":int}
        self.last_signal = {}  # code -> last signal type

    def reset(self, code: str = None):
        if code is None:
            self._bar_state.clear()
            self._closes.clear()
            self._counts.clear()
            self.last_signal.clear()
        else:
            self._bar_state.pop(code, None)
            self._closes.pop(code, None)
            self._counts.pop(code, None)
            self.last_signal.pop(code, None)

    def _log(self, code: str, message: str):
        if self.verbose:
            ##print(f"【同花顺九转】{code} {message}")
            pass

    def _is_trading_time(self, dt: datetime) -> bool:
        if not self.restrict_session:
            return True

        tm = dt.time()
        in_morning = (tm >= datetime.strptime("09:30:00", "%H:%M:%S").time() and
                      tm <= datetime.strptime("11:30:00", "%H:%M:%S").time())
        in_afternoon = (tm >= datetime.strptime("13:00:00", "%H:%M:%S").time() and
                        tm <= datetime.strptime("15:00:00", "%H:%M:%S").time())
        return in_morning or in_afternoon

    def on_tick(self, code: str, timestamp, price) -> str:
        if price is None:
            return None

        try:
            if isinstance(timestamp, (int, float)):
                ts_sec = timestamp if timestamp < 1e10 else timestamp / 1000
                dt = datetime.fromtimestamp(ts_sec)
            else:
                dt = pd.to_datetime(timestamp)
        except Exception:
            return None

        if not self._is_trading_time(dt):
            return None

        minute_bucket = int(dt.timestamp() // 60)
        bar = self._bar_state.get(code)

        cnt = self._counts.setdefault(code, {"up": 0, "down": 0})
        closes = self._closes.setdefault(code, deque(maxlen=1000))

        signal = None

        if bar is None or minute_bucket != bar["minute"]:
            if bar is not None:
                closes.append(bar["close"])
                signal = self._check_signal(code, closes, cnt)

            self._bar_state[code] = {
                "minute": minute_bucket,
                "open": float(price),
                "high": float(price),
                "low": float(price),
                "close": float(price)
            }
        else:
            bar["close"] = float(price)
            if price > bar["high"]:
                bar["high"] = float(price)
            if price < bar["low"]:
                bar["low"] = float(price)

        return signal

    def _check_signal(self, code: str, closes: deque, cnt: dict) -> str:
        if len(closes) <= self.compare_offset:
            return None

        cur = closes[-1]
        ref = closes[-1 - self.compare_offset]

        if cur >= ref:
            cnt["up"] += 1
        else:
            cnt["up"] = 0

        if cur <= ref:
            cnt["down"] += 1
        else:
            cnt["down"] = 0

        signal = None
        if cnt["up"] >= self.required_bars:
            signal = "卖出"
            cnt["up"] = 0
            self._log(code, f"触发卖出信号，连续{self.required_bars}次收盘价不低于前{self.compare_offset}根K线")
        elif cnt["down"] >= self.required_bars:
            signal = "买入"
            cnt["down"] = 0
            self._log(code, f"触发买入信号，连续{self.required_bars}次收盘价不高于前{self.compare_offset}根K线")

        if signal:
            self.last_signal[code] = signal
        return signal

    def warmup_from_history(self, code: str, history_data: pd.DataFrame):
        if not isinstance(history_data.index, pd.DatetimeIndex):
            raise ValueError("历史数据必须有DatetimeIndex")

        df = history_data.resample('1min').agg({
            'price': 'ohlc'
        })
        df.columns = df.columns.droplevel(0)

        for _, row in df.iterrows():
            if pd.isna(row['close']):
                continue
            self.on_tick(code, row.name, row['close'])


